#include <torch/torch.h>

#include <chrono>
#include <filesystem>
#include <iostream>
#include <string>
#include <vector>

#include "DoseStatTorch.h"
#include "MyMacro.h"
#include "cnpy.h"

///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 控制是否打印
#if 0
#    define MYDBG MYINFO
#else
#    define MYDBG(...)
#endif

///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 测试 dose stat torch
static int tt()
{
    namespace fs = std::filesystem;

#if 0  // in-place 测试数据
        std::vector<double> voxelVolume = {1.0, 2.0, 3.0, 4.0, 5.0};
        std::vector<double> voxelDose = {10.0, 20.0, 30.0, 40.0, 50.0};
        bool doExactAccumulation = true;
#endif

    // 从npz文件读取数据
    std::vector<double> voxelVolume;
    std::vector<double> voxelDose;
    bool doExactAccumulation = false;

    try
    {
#if 0
            const auto& dsInfo = args()->GetTe_DoseStat();
            const auto& fPath = fs::path(dsInfo.filePath).generic_string();

            // 对应的 "prescription_type" : "D_95"
            cnpy::npz_t npz_datas = cnpy::npz_load(fPath);
            for (const auto& [key, value] : npz_datas)
            {
                MYINFO("array key: {}", key);
            }
            cnpy::NpyArray arr_volumes = npz_datas[dsInfo.volumeName];  // 数据存储在键 "volumes"下
            auto* volumes_data = arr_volumes.data<double>();
            voxelVolume = std::vector<double>(volumes_data, volumes_data + arr_volumes.num_vals);
            // 读取剂量数据
            cnpy::NpyArray arr_doses = npz_datas[dsInfo.doseName];  // 假设数据存储在键"doses"下
            auto* doses_data = arr_doses.data<double>();
            voxelDose = std::vector<double>(doses_data, doses_data + arr_doses.num_vals);
#endif
        const auto& fPath = fs::path(MY_TESTIO_CASE1).generic_string();

        // 对应的 "prescription_type" : "D_95"
        cnpy::npz_t npz_datas = cnpy::npz_load(fPath);
        for (const auto& [key, value] : npz_datas)
        {
            MYINFO("array key: {}", key);
        }
        cnpy::NpyArray arr_volumes = npz_datas["volumes"];  // 数据存储在键 "volumes"下
        auto* volumes_data = arr_volumes.data<double>();
        voxelVolume = std::vector<double>(volumes_data, volumes_data + arr_volumes.num_vals);
        // 读取剂量数据
        cnpy::NpyArray arr_doses = npz_datas["doses"];  // 假设数据存储在键 "doses"下
        auto* doses_data = arr_doses.data<double>();
        voxelDose = std::vector<double>(doses_data, doses_data + arr_doses.num_vals);

        MYINFO("Success: read data from .npz: ");
        MYINFO("  voxelVolume count: {}", voxelVolume.size());
        MYINFO("  voxelDose count: {}", voxelDose.size());

        // 确保两个数组长度相同
        if (voxelVolume.size() != voxelDose.size())
        {
            MYERROR("Error: voxelVolume and voxelDose have different lengths!");
            // 取较短的长度
            size_t min_size = std::min(voxelVolume.size(), voxelDose.size());
            voxelVolume.resize(min_size);
            voxelDose.resize(min_size);
        }
    }
    catch (const std::exception& e)
    {
        MYERROR("read npz file failed: {}", e.what());
        // 使用默认测试数据
        voxelVolume = {1.0, 2.0, 3.0, 4.0, 5.0};
        voxelDose = {10.0, 20.0, 30.0, 40.0, 50.0};
    }

    MYINFO("Testing Dose Statistics Functions");
    MYINFO("==================================");

    // 测试accumulate函数
    MYINFO("\n1. Testing Accumulate function...");

    auto volumeTc = torch::tensor(voxelVolume, torch::kFloat64);
    auto doseTc = torch::tensor(voxelDose, torch::kFloat64);

    auto [binDoseTc, accRelVolTc] = dose_torch::accumulate(volumeTc, doseTc, doExactAccumulation);

#if 0
    MYINFO("   Binned Dose Values: ");
    const auto dose_p = binDoseTc.data_ptr<double>();
    for (size_t i = 0; i < binDoseTc.size(0); ++i)
    {
        MYINFO("{} ", dose_p[i]);
    }
    MYINFO("");

    MYINFO("   Accumulated Relative Volumes: ");
    const auto volume_p = accRelVolTc.data_ptr<double>();
    for (size_t i = 0; i < accRelVolTc.size(0); ++i)
    {
        MYINFO("{} ", volume_p[i]);
    }
    MYINFO("");
#endif
    // 测试get_metric函数
    MYINFO("\n2. Testing GetMetric function...");

    double d98 = dose_torch::get_metric(binDoseTc, accRelVolTc, 0.98).toDouble();
    MYINFO("   D_98 = {} Gy", d98);

    double d95 = dose_torch::get_metric(binDoseTc, accRelVolTc, 0.95).toDouble();
    MYINFO("   D_95 = {} Gy", d95);

    double d50 = dose_torch::get_metric(binDoseTc, accRelVolTc, 0.5).toDouble();
    MYINFO("   D_50 = {} Gy", d50);

    // 测试get_dose_statistics_roi_interp函数
    MYINFO("\n3. Testing GetDoseStatRoiInterp function...");

    // 测试平均剂量
    double averageDose = dose_torch::get_dose_stat_roi_interp(binDoseTc, accRelVolTc, volumeTc, doseTc, "average", 0.0).toDouble();
    MYINFO("   Average Dose = {} Gy", averageDose);

    // 测试D_95
    double d95_alt = dose_torch::get_dose_stat_roi_interp(binDoseTc, accRelVolTc, volumeTc, doseTc, "D_95", 0.0).toDouble();
    MYINFO("   D_95 (via interp) = {} Gy", d95_alt);

    // 测试自定义DoseAtVolume
    double doseAt70 = dose_torch::get_dose_stat_roi_interp(binDoseTc, accRelVolTc, volumeTc, doseTc, "DoseAtVolume", 0.7).toDouble();
    MYINFO("   Dose at 70% Volume = {} Gy", doseAt70);

    return 0;
}

///xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 测试 dose stat torch, 统计运行时间
static constexpr int PROFILE_LOOPN_ACCUM = 1000;   // 多次运行取平均值
static constexpr int PROFILE_LOOPN_METRIC = 1000;  // 多次运行取平均值
static constexpr int PROFILE_LOOPN_INTERP = 1000;  // 多次运行取平均值
static int tt_profile()
{
    namespace fs = std::filesystem;
    using namespace std::chrono;

    // 从npz文件读取数据
    std::vector<double> voxelVolume;
    std::vector<double> voxelDose;
    bool doExactAccumulation = false;

    try
    {
        const auto& fPath = fs::path(MY_TESTIO_CASE1).generic_string();

        // 对应的 "prescription_type" : "D_95"
        cnpy::npz_t npz_datas = cnpy::npz_load(fPath);
        for (const auto& [key, value] : npz_datas)
        {
            MYINFO("array key: {}", key);
        }
        cnpy::NpyArray arr_volumes = npz_datas["volumes"];  // 数据存储在键 "volumes"下
        auto* volumes_data = arr_volumes.data<double>();
        voxelVolume = std::vector<double>(volumes_data, volumes_data + arr_volumes.num_vals);
        // 读取剂量数据
        cnpy::NpyArray arr_doses = npz_datas["doses"];  // 假设数据存储在键 "doses"下
        auto* doses_data = arr_doses.data<double>();
        voxelDose = std::vector<double>(doses_data, doses_data + arr_doses.num_vals);

        MYINFO("Success: read data from .npz: ");
        MYINFO("  voxelVolume count: {}", voxelVolume.size());
        MYINFO("  voxelDose count: {}", voxelDose.size());

        // 确保两个数组长度相同
        if (voxelVolume.size() != voxelDose.size())
        {
            MYERROR("Error: voxelVolume and voxelDose have different lengths!");
            // 取较短的长度
            size_t min_size = std::min(voxelVolume.size(), voxelDose.size());
            voxelVolume.resize(min_size);
            voxelDose.resize(min_size);
        }
    }
    catch (const std::exception& e)
    {
        MYERROR("read npz file failed: {}", e.what());
        std::exit(EXIT_FAILURE);
    }

    MYINFO("Testing Dose Statistics Functions");
    MYINFO("==================================");

    ///=============================================================================== 测试accumulate函数
    MYINFO("\n1. Testing Accumulate function...");
    auto volumeTc = torch::tensor(voxelVolume, torch::kFloat64);
    auto doseTc = torch::tensor(voxelDose, torch::kFloat64);
    ///========================================== time profile
    {
        const auto start = high_resolution_clock::now();
        for (int i = 0; i < PROFILE_LOOPN_ACCUM; i++)
        {
            auto [binDoseTc, accRelVolTc] = dose_torch::accumulate(volumeTc, doseTc, doExactAccumulation);
        }
        const auto end = high_resolution_clock::now();
        const auto totalTimeAccum = duration_cast<microseconds>(end - start).count();
        const double avgTimeAccum = double(totalTimeAccum) / double(PROFILE_LOOPN_ACCUM);
        MYINFO("\taccumulate: average execution time ({} runs): {:.2f} microseconds", PROFILE_LOOPN_ACCUM, avgTimeAccum);
    }

    ///=============================================================================== 测试get_metric函数
    MYDBG("\n2. Testing GetMetric function...");
    auto [binDoseTc, accRelVolTc] = dose_torch::accumulate(volumeTc, doseTc, doExactAccumulation);
    double totalMetricTime = 0.0;
    for (int i = 0; i < PROFILE_LOOPN_METRIC; i++)
    {
        const auto start = high_resolution_clock::now();

        double d98 = dose_torch::get_metric(binDoseTc, accRelVolTc, 0.98).toDouble();
        MYDBG("   D_98 = {} Gy", d98);
        double d95 = dose_torch::get_metric(binDoseTc, accRelVolTc, 0.95).toDouble();
        MYDBG("   D_95 = {} Gy", d95);
        double d50 = dose_torch::get_metric(binDoseTc, accRelVolTc, 0.5).toDouble();
        MYDBG("   D_50 = {} Gy", d50);

        const auto end = high_resolution_clock::now();
        const auto duration = duration_cast<microseconds>(end - start).count();
        totalMetricTime += duration;
    }
    const double avgMetricTime = totalMetricTime / PROFILE_LOOPN_METRIC;
    MYINFO("\tget_metric: average execution time ({} runs): {:.2f} microseconds", PROFILE_LOOPN_METRIC, avgMetricTime);

    ///=============================================================================== 测试get_dose_statistics_roi_interp函数
    MYDBG("\n3. Testing GetDoseStatRoiInterp function...");
    double totalInterpTime = 0.0;
    double averageDose = 0.0, d95_alt = 0.0, doseAt70 = 0.0;
    for (int i = 0; i < PROFILE_LOOPN_INTERP; i++)
    {
        const auto start = high_resolution_clock::now();

        // 测试平均剂量
        averageDose = dose_torch::get_dose_stat_roi_interp(binDoseTc, accRelVolTc, volumeTc, doseTc, "average", 0.0).toDouble();
        MYDBG("   Average Dose = {} Gy", averageDose);
        // 测试D_95
        d95_alt = dose_torch::get_dose_stat_roi_interp(binDoseTc, accRelVolTc, volumeTc, doseTc, "D_95", 0.0).toDouble();
        MYDBG("   D_95 (via interp) = {} Gy", d95_alt);
        // 测试自定义DoseAtVolume
        doseAt70 = dose_torch::get_dose_stat_roi_interp(binDoseTc, accRelVolTc, volumeTc, doseTc, "DoseAtVolume", 0.7).toDouble();
        MYDBG("   Dose at 70% Volume = {} Gy", doseAt70);

        const auto end = high_resolution_clock::now();
        const auto duration = duration_cast<microseconds>(end - start).count();
        totalInterpTime += duration;
    }

    const double avgInterpTime = totalInterpTime / double(PROFILE_LOOPN_INTERP);
    MYINFO("\tAverage Dose = {} Gy", averageDose);
    MYINFO("\tD_95 (via interp) = {} Gy", d95_alt);
    MYINFO("\tDose at 70% Volume = {} Gy", doseAt70);
    MYINFO("\tget_dose_stat_roi_interp: average execution time ({} runs): {:.2f} microseconds", PROFILE_LOOPN_INTERP, avgInterpTime);

    return 0;
}

int main()
{
    MYINFO("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx tt()");
    tt();

    MYINFO("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx tt_profile()");
    tt_profile();

    return 0;
}